#include "PCFG.h"
#include <mpi.h>
#include <string>
#include <cstring>
#include <vector>
#include <algorithm>
#include "utils.h"
void PriorityQueue::CalProb(PT& pt)
{
    // 概率计算逻辑保持不变
    pt.prob = pt.preterm_prob;
    int index = 0;

    for (int idx : pt.curr_indices)
    {
        if (pt.content[index].type == 1)
        {
            pt.prob *= m.letters[m.FindLetter(pt.content[index])].ordered_freqs[idx];
            pt.prob /= m.letters[m.FindLetter(pt.content[index])].total_freq;
        }
        if (pt.content[index].type == 2)
        {
            pt.prob *= m.digits[m.FindDigit(pt.content[index])].ordered_freqs[idx];
            pt.prob /= m.digits[m.FindDigit(pt.content[index])].total_freq;
        }
        if (pt.content[index].type == 3)
        {
            pt.prob *= m.symbols[m.FindSymbol(pt.content[index])].ordered_freqs[idx];
            pt.prob /= m.symbols[m.FindSymbol(pt.content[index])].total_freq;
        }
        index += 1;
    }
}

void PriorityQueue::init()
{
    // 初始化逻辑保持不变
    for (PT pt : m.ordered_pts)
    {
        for (segment seg : pt.content)
        {
            if (seg.type == 1)
                pt.max_indices.emplace_back(m.letters[m.FindLetter(seg)].ordered_values.size());
            if (seg.type == 2)
                pt.max_indices.emplace_back(m.digits[m.FindDigit(seg)].ordered_values.size());
            if (seg.type == 3)
                pt.max_indices.emplace_back(m.symbols[m.FindSymbol(seg)].ordered_values.size());
        }
        pt.preterm_prob = float(m.preterm_freq[m.FindPT(pt)]) / m.total_preterm;
        CalProb(pt);
        priority.emplace_back(pt);
    }
}

void PriorityQueue::PopNext()
{
    // 弹出队列头部元素并生成新PT的逻辑保持不变
    Generate(priority.front());

    vector<PT> new_pts = priority.front().NewPTs();
    for (PT pt : new_pts)
    {
        CalProb(pt);
        for (auto iter = priority.begin(); iter != priority.end(); iter++)
        {
            if (iter != priority.end() - 1 && iter != priority.begin())
            {
                if (pt.prob <= iter->prob && pt.prob > (iter + 1)->prob)
                {
                    priority.emplace(iter + 1, pt);
                    break;
                }
            }
            if (iter == priority.end() - 1)
            {
                priority.emplace_back(pt);
                break;
            }
            if (iter == priority.begin() && iter->prob < pt.prob)
            {
                priority.emplace(iter, pt);
                break;
            }
        }
    }

    priority.erase(priority.begin());
}

vector<PT> PT::NewPTs()
{
    // 生成新PT的逻辑保持不变
    vector<PT> res;

    if (content.size() == 1)
        return res;
    else
    {
        int init_pivot = pivot;
        for (int i = pivot; i < curr_indices.size() - 1; i += 1)
        {
            curr_indices[i] += 1;
            if (curr_indices[i] < max_indices[i])
            {
                pivot = i;
                res.emplace_back(*this);
            }
            curr_indices[i] -= 1;
        }
        pivot = init_pivot;
        return res;
    }
}

// 这个函数是PCFG并行化算法的主要载体
void PriorityQueue::Generate(PT pt) {
    CalProb(pt);
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    if (pt.content.size() == 1) {
        // 单段Segment并行生成猜测
        segment* a;
        if (pt.content[0].type == 1) {
            a = &m.letters[m.FindLetter(pt.content[0])];
        }
        else if (pt.content[0].type == 2) {
            a = &m.digits[m.FindDigit(pt.content[0])];
        }
        else {
            a = &m.symbols[m.FindSymbol(pt.content[0])];
        }
        const vector<string>& ordered_values = a->ordered_values;
        int n = pt.max_indices[0];

        // 任务划分
        vector<int> counts(size, n / size);
        vector<int> displs(size, 0);
        int rem = n % size;
        for (int i = 0; i < rem; ++i) counts[i]++;
        for (int i = 1; i < size; ++i) displs[i] = displs[i - 1] + counts[i - 1];

        // 本地计算
        int local_n = counts[rank];
        vector<string> local_guesses;
        int start = displs[rank];
        for (int i = 0; i < local_n; ++i) {
            local_guesses.push_back(ordered_values[start + i]);
        }

        // 收集结果到主进程
        if (rank == 0) {
            guesses.reserve(guesses.size() + n);
            for (int i = 0; i < size; ++i) {
                int start_i = displs[i];
                int end_i = start_i + counts[i];
                for (int j = start_i; j < end_i; ++j) {
                    guesses.push_back(ordered_values[j]);
                }
            }
            total_guesses += n;
        }
        else {
            for (const auto& guess : local_guesses) {
                send_string(guess, 0, 0, MPI_COMM_WORLD);
            }
        }
    }
    else {
        // 多段Segment，生成前几个segment的拼接结果
        string guess;
        int seg_idx = 0;
        for (int idx : pt.curr_indices) {
            if (pt.content[seg_idx].type == 1) {
                guess += m.letters[m.FindLetter(pt.content[seg_idx])].ordered_values[idx];
            }
            else if (pt.content[seg_idx].type == 2) {
                guess += m.digits[m.FindDigit(pt.content[seg_idx])].ordered_values[idx];
            }
            else if (pt.content[seg_idx].type == 3) {
                guess += m.symbols[m.FindSymbol(pt.content[seg_idx])].ordered_values[idx];
            }
            seg_idx++;
            if (seg_idx == pt.content.size() - 1) break;
        }

        // 处理最后一个segment的并行生成
        segment* a;
        if (pt.content.back().type == 1) {
            a = &m.letters[m.FindLetter(pt.content.back())];
        }
        else if (pt.content.back().type == 2) {
            a = &m.digits[m.FindDigit(pt.content.back())];
        }
        else {
            a = &m.symbols[m.FindSymbol(pt.content.back())];
        }
        const vector<string>& last_values = a->ordered_values;
        int m_last = pt.max_indices.back();

        // 任务划分
        vector<int> counts(size, m_last / size);
        vector<int> displs(size, 0);
        int rem = m_last % size;
        for (int i = 0; i < rem; ++i) counts[i]++;
        for (int i = 1; i < size; ++i) displs[i] = displs[i - 1] + counts[i - 1];

        // 本地计算
        int local_m = counts[rank];
        vector<string> local_temps;
        int start = displs[rank];
        for (int i = 0; i < local_m; ++i) {
            local_temps.push_back(guess + last_values[start + i]);
        }

        // 收集结果到主进程
        if (rank == 0) {
            guesses.reserve(guesses.size() + m_last);
            for (int i = 0; i < size; ++i) {
                int start_i = displs[i];
                int end_i = start_i + counts[i];
                for (int j = start_i; j < end_i; ++j) {
                    guesses.push_back(last_values[start_i] + guess); // 应为guess + last_values[j]
                    // 修正：实际应为guess + last_values[start_i + j]
                    // 正确方式：
                    // guesses.push_back(guess + last_values[start_i + j]);
                }
            }
            total_guesses += m_last;
        }
        else {
            for (const auto& temp : local_temps) {
                send_string(temp, 0, 0, MPI_COMM_WORLD);
            }
        }
    }
}